"""
AI-based mesh quality prediction module
"""
import numpy as np
import os
import pickle
from pathlib import Path

class MeshQualityPredictor:
    """Class for predicting mesh quality using AI models"""
    
    def __init__(self, model_path=None):
        """
        Initialize the mesh quality predictor
        
        Parameters:
        -----------
        model_path : str, optional
            Path to the pre-trained model file
        """
        self.model = None
        self.model_path = model_path
        
        if model_path and os.path.exists(model_path):
            self.load_model(model_path)
    
    def load_model(self, model_path):
        """
        Load a pre-trained model
        
        Parameters:
        -----------
        model_path : str
            Path to the model file
        """
        try:
            with open(model_path, 'rb') as f:
                self.model = pickle.load(f)
            print(f"Model loaded from {model_path}")
        except Exception as e:
            print(f"Error loading model: {e}")
    
    def save_model(self, model_path=None):
        """
        Save the trained model
        
        Parameters:
        -----------
        model_path : str, optional
            Path to save the model file
        """
        if self.model is None:
            print("No model to save")
            return
        
        if model_path is None:
            model_path = self.model_path
        
        if model_path is None:
            model_path = os.path.join(
                Path(__file__).parent.parent.parent, 
                'data', 
                'quality_predictor_model.pkl'
            )
        
        # Ensure directory exists
        os.makedirs(os.path.dirname(model_path), exist_ok=True)
        
        try:
            with open(model_path, 'wb') as f:
                pickle.dump(self.model, f)
            print(f"Model saved to {model_path}")
        except Exception as e:
            print(f"Error saving model: {e}")
    
    def extract_features(self, mesh_data):
        """
        Extract features from mesh data for quality prediction
        
        Parameters:
        -----------
        mesh_data : dict
            Dictionary containing mesh data
            
        Returns:
        --------
        numpy.ndarray
            Array of features for each element
        """
        points = mesh_data['points']
        cells = mesh_data['cells']
        
        features = []
        
        for cell in cells:
            if len(cell) == 3:  # Triangle
                p1, p2, p3 = points[cell[0]], points[cell[1]], points[cell[2]]
                
                # Calculate edge lengths
                a = np.linalg.norm(p2 - p1)
                b = np.linalg.norm(p3 - p2)
                c = np.linalg.norm(p1 - p3)
                
                # Calculate angles
                angle_a = np.arccos(np.clip(np.dot(p2-p1, p3-p1) / (a * c), -1.0, 1.0))
                angle_b = np.arccos(np.clip(np.dot(p1-p2, p3-p2) / (a * b), -1.0, 1.0))
                angle_c = np.arccos(np.clip(np.dot(p1-p3, p2-p3) / (b * c), -1.0, 1.0))
                
                # Calculate area
                s = (a + b + c) / 2
                area = np.sqrt(max(0, s * (s - a) * (s - b) * (s - c)))
                
                # Extract features
                cell_features = [
                    a, b, c,  # Edge lengths
                    angle_a, angle_b, angle_c,  # Angles
                    area,  # Area
                    max(a, b, c) / min(a, b, c) if min(a, b, c) > 0 else 100,  # Edge ratio
                    min(angle_a, angle_b, angle_c),  # Minimum angle
                    max(angle_a, angle_b, angle_c)  # Maximum angle
                ]
                
                features.append(cell_features)
        
        return np.array(features)
    
    def train(self, mesh_data_list, quality_values_list):
        """
        Train the quality prediction model
        
        Parameters:
        -----------
        mesh_data_list : list
            List of mesh data dictionaries
        quality_values_list : list
            List of quality values for each mesh
            
        Returns:
        --------
        bool
            True if training was successful, False otherwise
        """
        try:
            # Import scikit-learn (if available)
            from sklearn.ensemble import RandomForestRegressor
            from sklearn.model_selection import train_test_split
            from sklearn.metrics import mean_squared_error
        except ImportError:
            print("scikit-learn is not available. Install it with 'pip install scikit-learn'")
            return False
        
        # Extract features from all meshes
        X = []
        y = []
        
        for mesh_data, quality_values in zip(mesh_data_list, quality_values_list):
            features = self.extract_features(mesh_data)
            X.append(features)
            y.append(quality_values)
        
        # Flatten the data
        X_flat = np.vstack(X)
        y_flat = np.concatenate(y)
        
        # Split data into training and validation sets
        X_train, X_val, y_train, y_val = train_test_split(
            X_flat, y_flat, test_size=0.2, random_state=42
        )
        
        # Train a random forest regressor
        model = RandomForestRegressor(
            n_estimators=100, 
            max_depth=10,
            random_state=42
        )
        
        model.fit(X_train, y_train)
        
        # Evaluate the model
        y_pred = model.predict(X_val)
        mse = mean_squared_error(y_val, y_pred)
        print(f"Model validation MSE: {mse:.4f}")
        
        # Save the model
        self.model = model
        
        return True
    
    def predict(self, mesh_data):
        """
        Predict quality values for a mesh
        
        Parameters:
        -----------
        mesh_data : dict
            Dictionary containing mesh data
            
        Returns:
        --------
        numpy.ndarray
            Predicted quality values for each element
        """
        if self.model is None:
            print("No model loaded. Using fallback quality estimation.")
            # Fallback to a simple heuristic
            return self._fallback_quality_estimation(mesh_data)
        
        # Extract features
        features = self.extract_features(mesh_data)
        
        # Make predictions
        predictions = self.model.predict(features)
        
        return predictions
    
    def _fallback_quality_estimation(self, mesh_data):
        """
        Fallback quality estimation when no model is available
        
        This uses a simple heuristic based on aspect ratio.
        """
        points = mesh_data['points']
        cells = mesh_data['cells']
        
        quality_values = []
        
        for cell in cells:
            if len(cell) == 3:  # Triangle
                p1, p2, p3 = points[cell[0]], points[cell[1]], points[cell[2]]
                
                # Calculate edge lengths
                a = np.linalg.norm(p2 - p1)
                b = np.linalg.norm(p3 - p2)
                c = np.linalg.norm(p1 - p3)
                
                # Calculate semi-perimeter
                s = (a + b + c) / 2
                
                # Calculate area using Heron's formula
                area = np.sqrt(max(0, s * (s - a) * (s - b) * (s - c)))
                
                # Calculate aspect ratio (circumradius to inradius ratio)
                if area > 1e-10:  # Avoid division by zero
                    inradius = area / s
                    circumradius = (a * b * c) / (4 * area)
                    aspect_ratio = circumradius / inradius
                else:
                    aspect_ratio = 100.0
                
                quality_values.append(aspect_ratio)
        
        return np.array(quality_values) 